package com.hht.toolanduilibs.util.usbcamera

import android.graphics.*
import android.media.Image
import android.media.ImageReader
import android.os.Build
import java.io.ByteArrayOutputStream
import java.nio.ByteBuffer


/**
 * Author:chen'hu
 * Time:2024/5/14 15:01
 * Description:This is ImageToYuvHelps
 */
object ImageToYuvHelps {
    private const val YUV420P = 0
    private const val YUV420SP = 1
    private const val NV21 = 2

    fun readMJPEG(reader: ImageReader): ByteArray? {
        var mjpegData: ByteArray? = null

        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.KITKAT) {
            return mjpegData
        }
        var image: Image? = reader.acquireLatestImage()

        if (image != null && image?.planes != null) {
            val buffer = image.planes[0].buffer
            mjpegData = ByteArray(buffer.remaining())
            //TODO 把buffer数据写入到data中
            buffer.get(mjpegData)
        }
        image?.close()

        return mjpegData
    }


    fun readImageReaderYuv(reader: ImageReader, type: Int = NV21): ByteArray? {
        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.KITKAT) {
            return null
        }
        var image: Image? = reader.acquireLatestImage() ?: return null
        if (image == null || image.planes == null || image.planes.isEmpty()) return null
        //数据有效宽度，一般的，图片width <= rowStride，这也是导致byte[].length <= capacity的原因
        // 所以我们只取width部分
        //获取源数据，如果是YUV格式的数据planes.length = 3
        //plane[i]里面的实际数据可能存在byte[].length <= capacity (缓冲区总大小)
        val planes = image.planes
        //数据有效宽度，一般的，图片width <= rowStride，这也是导致byte[].length <= capacity的原因
        // 所以我们只取width部分
        val width = image.width
        val height = image.height
        //此处用来装填最终的YUV数据，需要1.5倍的图片大小，因为Y U V 比例为 4:1:1
        val yuvBytes =
            ByteArray(width * height * ImageFormat.getBitsPerPixel(ImageFormat.YUV_420_888) / 8)
        //目标数组的装填到的位置
        var dstIndex = 0

        //临时存储uv数据的
        val uBytes = ByteArray(width * height / 4)
        val vBytes = ByteArray(width * height / 4)
        var uIndex = 0
        var vIndex = 0

        var pixelsStride: Int
        var rowStride: Int
        for (i in planes.indices) {
            pixelsStride = planes[i].pixelStride
            rowStride = planes[i].rowStride
            val buffer: ByteBuffer = planes[i].buffer
            //如果pixelsStride==2，一般的Y的buffer长度=640*480，UV的长度=640*480/2-1
            //源数据的索引，y的数据是byte中连续的，u的数据是v向左移以为生成的，两者都是偶数位为有效数据
            val bytes = ByteArray(buffer.capacity())
            buffer[bytes]
            var srcIndex = 0
            when (i) {
                0 -> {
                    //直接取出来所有Y的有效区域，也可以存储成一个临时的bytes，到下一步再copy
                    for (j in 0 until height) {
                        System.arraycopy(bytes, srcIndex, yuvBytes, dstIndex, width)
                        srcIndex += rowStride
                        dstIndex += width
                    }
                }
                1 -> {
                    //根据pixelsStride取相应的数据
                    for (j in 0 until height / 2) {
                        for (k in 0 until width / 2) {
                            uBytes[uIndex++] = bytes[srcIndex]
                            srcIndex += pixelsStride
                        }
                        if (pixelsStride == 2) {
                            srcIndex += rowStride - width
                        } else if (pixelsStride == 1) {
                            srcIndex += rowStride - width / 2
                        }
                    }
                }
                2 -> {
                    //根据pixelsStride取相应的数据
                    for (j in 0 until height / 2) {
                        for (k in 0 until width / 2) {
                            vBytes[vIndex++] = bytes[srcIndex]
                            srcIndex += pixelsStride
                        }
                        if (pixelsStride == 2) {
                            srcIndex += rowStride - width
                        } else if (pixelsStride == 1) {
                            srcIndex += rowStride - width / 2
                        }
                    }
                }
            }
        }
        //根据要求的结果类型进行填充
        when (type) {
            YUV420P -> {
                System.arraycopy(uBytes, 0, yuvBytes, dstIndex, uBytes.size)
                System.arraycopy(vBytes, 0, yuvBytes, dstIndex + uBytes.size, vBytes.size)
            }
            YUV420SP -> {
                var i = 0
                while (i < vBytes.size) {
                    yuvBytes[dstIndex++] = uBytes[i]
                    yuvBytes[dstIndex++] = vBytes[i]
                    i++
                }
            }
            NV21 -> {
                var i = 0
                while (i < vBytes.size) {
                    yuvBytes[dstIndex++] = vBytes[i]
                    yuvBytes[dstIndex++] = uBytes[i]
                    i++
                }
            }
        }
        image.close()
        return yuvBytes
    }

    /***
     * YUV420 转化成 RGB
     */
    fun decodeYuvToRgb(yuv420sp: ByteArray, width: Int, height: Int): IntArray {
        val frameSize = width * height
        val rgb = IntArray(frameSize)
        var j = 0
        var yp = 0
        while (j < height) {
            var uvp = frameSize + (j shr 1) * width
            var u = 0
            var v = 0
            var i = 0
            while (i < width) {
                var y = (0xff and yuv420sp[yp].toInt()) - 16
                if (y < 0) y = 0
                if (i and 1 == 0) {
                    v = (0xff and yuv420sp[uvp++].toInt()) - 128
                    u = (0xff and yuv420sp[uvp++].toInt()) - 128
                }
                val y1192 = 1192 * y
                var r = y1192 + 1634 * v
                var g = y1192 - 833 * v - 400 * u
                var b = y1192 + 2066 * u
                if (r < 0) r = 0 else if (r > 262143) r = 262143
                if (g < 0) g = 0 else if (g > 262143) g = 262143
                if (b < 0) b = 0 else if (b > 262143) b = 262143
                rgb[yp] = (-0x1000000 or (r shl 6 and 0xff0000)
                        or (g shr 2 and 0xff00) or (b shr 10 and 0xff))
                i++
                yp++
            }
            j++
        }
        return rgb
    }

    fun readNV21Data(reader: ImageReader): ByteArray? {
        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.KITKAT) {
            return null
        }
        var image: Image? = reader.acquireLatestImage() ?: return null
        if (image == null || image.planes == null || image.planes.isEmpty()) return null
        var nv21Data: ByteArray? = null
        try {
            val crop = image.cropRect
            val format = image.format
            val width = crop.width()
            val height = crop.height()
            val planes = image.planes
            nv21Data = ByteArray(width * height * ImageFormat.getBitsPerPixel(format) / 8)
            val rowData = ByteArray(planes[0].rowStride)
            var channelOffset = 0
            var outputStride = 1
            for (i in planes.indices) {
                when (i) {
                    0 -> {
                        channelOffset = 0
                        outputStride = 1
                    }
                    1 -> {
                        channelOffset = width * height + 1
                        outputStride = 2
                    }
                    2 -> {
                        channelOffset = width * height
                        outputStride = 2
                    }
                }
                val buffer = planes[i].buffer
                val rowStride = planes[i].rowStride
                val pixelStride = planes[i].pixelStride
                val shift = if (i == 0) 0 else 1
                val w = width shr shift
                val h = height shr shift
                buffer.position(rowStride * (crop.top shr shift) + pixelStride * (crop.left shr shift))
                for (row in 0 until h) {
                    var length: Int
                    if (pixelStride == 1 && outputStride == 1) {
                        length = w
                        buffer[nv21Data, channelOffset, length]
                        channelOffset += length
                    } else {
                        length = (w - 1) * pixelStride + 1
                        buffer[rowData, 0, length]
                        for (col in 0 until w) {
                            nv21Data[channelOffset] = rowData[col * pixelStride]
                            channelOffset += outputStride
                        }
                    }
                    if (row < h - 1) {
                        buffer.position(buffer.position() + rowStride - length)
                    }
                }
            }
        } catch (e: Exception) {
            e.printStackTrace()
        } finally {
            image.close()
        }
        return nv21Data
    }

    fun nv21ToBitmap(nv21: ByteArray, width: Int, height: Int): Bitmap {
        val yuvImage = YuvImage(nv21, ImageFormat.NV21, width, height, null)
        val out = ByteArrayOutputStream(nv21.size)
        yuvImage.compressToJpeg(Rect(0, 0, width, height), 100, out)
        val imageBytes = out.toByteArray()
        return BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.size)
    }
}